import os
import json
import glob


def main():
    path = os.path.dirname(os.path.abspath(__file__))
    jsonl_list = glob.glob(os.path.join(path, "visualglm_sample_2w", "*.jsonl"))
    json_data_list = []
    for jsonl_file in jsonl_list:
        with open(jsonl_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                json_data = json.loads(line.strip())
                img = json_data['img']
                json_data['img'] = os.path.join("visualglm_sample_2w", img)
                if os.path.exists(os.path.join(path, "train2017", img)):
                    os.system(f"cp {os.path.join(path, 'train2017', img)} {os.path.join(path, 'visualglm_sample_2w')}")
                elif os.path.exists(os.path.join(path, "test2017", img)):
                    os.system(f"cp {os.path.join(path, 'test2017', img)} {os.path.join(path, 'visualglm_sample_2w')}")
                elif os.path.exists(os.path.join(path, "val2017", img)):
                    os.system(f"cp {os.path.join(path, 'val2017', img)} {os.path.join(path, 'visualglm_sample_2w')}")
                else:
                    continue
                json_data_list.append(json.dumps(json_data, ensure_ascii=False))
    with open(os.path.join(path, "visualglm_sample_2w", "dataset.jsonl"), 'w') as f:
        f.write('\n'.join(json_data_list))

if __name__ == "__main__":
    main()
